import subprocess
import re
import sys
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm
import argparse

def run_nmap_scan(ip, port):
    """执行Nmap扫描并返回结果"""
    try:
        command = ['nmap', '-Pn', '-p', str(port), ip]
        result = subprocess.run(
            command,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        return result.stdout
    except subprocess.CalledProcessError as e:
        print(f"扫描 {ip}:{port} 时出错: {e.stderr}", file=sys.stderr)
        return None
    except Exception as e:
        print(f"扫描 {ip}:{port} 时发生未知错误: {e}", file=sys.stderr)
        return None

def parse_nmap_output(output, ip, port):
    """解析Nmap输出，提取端口状态和服务信息"""
    if not output:
        return None
    
    # 改进的正则表达式，匹配更广泛的服务名称字符
    # 现在可以匹配包含连字符、点号等的服务名称
    pattern = re.compile(rf'\s*{port}/tcp\s+(\w+)\s+([\w\-.]+)')
    match = pattern.search(output)
    
    if match:
        state = match.group(1)
        service = match.group(2)
        return f"{ip}:{port} {port}/tcp {state} {service}"
    else:
        return f"{ip}:{port} 未找到匹配的端口信息"

def process_ip_port_pair(line):
    """处理单个IP:端口对"""
    line = line.strip()
    if not line:
        return None
    
    try:
        ip, port = line.split(':')
        port = int(port)
    except ValueError:
        print(f"无效的IP:端口格式: {line}", file=sys.stderr)
        return None
    
    output = run_nmap_scan(ip, port)
    return parse_nmap_output(output, ip, port)

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='批量Nmap端口扫描工具')
    parser.add_argument('-f', '--file', required=True, help='包含IP:端口列表的文件路径')
    parser.add_argument('-t', '--threads', type=int, default=4, help='并行线程数 (默认: 4)')
    parser.add_argument('-o', '--output', help='输出结果到文件')
    
    args = parser.parse_args()
    
    try:
        # 读取IP和端口列表
        with open(args.file, 'r') as f:
            lines = f.readlines()
        
        if not lines:
            print(f"错误: 文件 {args.file} 为空", file=sys.stderr)
            sys.exit(1)
        
        # 使用线程池并行处理
        with ThreadPoolExecutor(max_workers=args.threads) as executor:
            results = list(tqdm(
                executor.map(process_ip_port_pair, lines),
                total=len(lines),
                desc="扫描进度"
            ))
        
        # 过滤掉None结果并输出
        valid_results = [r for r in results if r]
        
        if args.output:
            with open(args.output, 'w') as f:
                f.write('\n'.join(valid_results))
            print(f"结果已保存到 {args.output}")
        else:
            print("\n扫描结果:")
            print('\n'.join(valid_results))
    
    except FileNotFoundError:
        print(f"错误: 找不到文件 {args.file}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"发生未知错误: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()